from transformers import AutoModelForCausalLM, AutoTokenizer
import json
import requests

# A2A机制: 在这个示例中，通过QueryAgent对ToolAgent的调用来体现A2A机制。这允许不同的智能体专注于各自的领域（如查询处理和工具执行），并通过定义好的接口相互通信。

# 加载预训练模型和分词器
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def generate_response(input_text, history=[]):
    # 根据历史对话构建上下文
    context = "\n".join(history) + "\n" if history else ""
    inputs = tokenizer.encode(context + input_text + tokenizer.eos_token, return_tensors='pt')
    
    outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

# ToolAgent: 该类实现了具体的工具功能，比如获取天气信息和进行简单计算。
class ToolAgent:
    """工具代理，执行特定任务"""
    def get_weather(self, location="Beijing"):
        response = requests.get("http://api.weatherapi.com/v1/current.json", params={"key": "your_api_key", "q": location})
        weather_data = response.json()
        return f"{location}当前温度是{weather_data['current']['temp_c']}度 Celsius。"

    def calculate(self, expr):
        try:
            result = eval(expr)
            return f"{expr} 的结果是 {result}"
        except Exception as e:
            return str(e)

# QueryAgent: 该类负责接收用户的输入，并根据需要调用ToolAgent的方法来获取更多信息。它还处理了从语言模型生成的响应，并在必要时解析出函数调用请求。
class QueryAgent:
    """查询代理，处理用户问题并决定是否需要调用ToolAgent"""
    def __init__(self, tool_agent: ToolAgent):
        self.tool_agent = tool_agent
    
    def handle_request(self, user_input, history=[]):
        response = generate_response(user_input, history)
        
        # 检查是否包含函数调用请求
        if 'function_call' in response:
            function_call = json.loads(response.split('function_call: ')[-1])
            function_name = function_call['name']
            arguments = function_call['arguments']

            # 调用ToolAgent的相应方法
            if function_name == "get_weather":
                execution_result = self.tool_agent.get_weather(arguments['location'])
            elif function_name == "calculate":
                execution_result = self.tool_agent.calculate(arguments['expression'])
            else:
                execution_result = "未知的功能调用。"
            
            final_response = f"{response}\n执行结果: {execution_result}"
        else:
            final_response = response
        
        # 更新历史记录
        history.append(f"用户: {user_input}")
        history.append(f"助手: {final_response}")
        return final_response

if __name__ == "__main__":
    tool_agent = ToolAgent()
    query_agent = QueryAgent(tool_agent)
    
    print("开始与智能助手对话（输入'exit'退出）：")
    history = []
    while True:
        input_text = input("你: ")
        if input_text.lower() == 'exit':
            break
        response = query_agent.handle_request(input_text, history)
        print(f"AI: {response}")